# -*- coding: UTF-8 -*-
import os
import numpy as np
import matplotlib.pyplot as plt
from keras.utils import Sequence
# from keras.preprocessing import sequence


# 计算振动数据标准差
def imsdatasigma(path, channel=0):
    filelist = os.listdir(path)
    num_of_files = len(filelist)
    datasigmalist = []
    for file in filelist:
        filepath = os.path.join(path, file)
        data = np.loadtxt(filepath)[:, channel]
        # sigma = np.sqrt(np.mean(np.square(data-data.mean())))
        sigma = np.std(data)
        datasigmalist.append(sigma)
    xlinespace = range(0, len(datasigmalist))
    plt.plot(xlinespace, datasigmalist)
    plt.xlabel('TimeStep')
    plt.ylabel('sigma')
    plt.savefig('sigma graph')
    plt.show()
    return datasigmalist


# 获取数据文件列表
def getIMSdatalist(path, step=3):
    filelist = os.listdir(path)
    num_of_files = len(filelist)
    assert num_of_files > step
    datalist = []
    for i in range(0, num_of_files-step+1) :
        dataitemlst = []
        for j in range(0, step):
            filepath = os.path.join(path, filelist[i+j])
            dataitemlst.append(filepath)
        dataitem = [dataitemlst, (i+step)/num_of_files]
        datalist.append(dataitem)
    # xlist = [row[0] for row in datalist]
    # ylist = [row[1] for row in datalist]
    return datalist


class IMSDataGenerator(Sequence):
    def __init__(self, datalist, chn=2, siglen=2000, batchsize=10, timestep=3, shuffle=True):
        self.datalist = datalist
        self.chn = chn
        self.siglen = siglen
        self.batchsize = batchsize
        self.shuffle = shuffle
        self.indexes = np.arange(len(self.datalist))
        self.timestep = timestep

    def __len__(self):
        return int(np.floor(len(self.datalist)/self.batchsize))

    def on_epoch_end(self):
        if self.shuffle == True:
            np.random.shuffle(self.indexes)

    def __getitem__(self, index):
        indexes = self.indexes[index*self.batchsize:(index+1)*self.batchsize]
        list_temp = [self.datalist[k] for k in indexes]
        x, y = self.__datageneration(list_temp)
        return x, y

    def __datageneration(self, indexes):
        x = []
        y = []
        for block in indexes:
            y.append(block[1])
            xtemp = []
            for item in block[0]:
                datablock = np.loadtxt(item)[:self.siglen, self.chn]
                xtemp.append(datablock)
            x.append(xtemp)
        # x = x.reshape((self.siglen,1))
        x = np.array(x)
        x = x.reshape(self.batchsize, self.timestep, self.siglen, 1)
        # x = np.reshape(x, (x.shape[:], 1))
        y = np.array(y)
        y = y.reshape(self.batchsize, 1)
        return x, y


def ims_dataset_split(datalist, shuffle=True, ratio=0.98):
    datasetidexes = np.arange(len(datalist))
    if shuffle == True:
        np.random.shuffle(datasetidexes)
    train_set_indexes = datasetidexes[0:int(len(datalist)*ratio)]
    test_set_indexes = datasetidexes[int(len(datalist)*ratio):]
    # train_set_list = [datalist[i] for i in train_set_indexes]
    # test_set_list = [datalist[k] for k in test_set_indexes ]
    return train_set_indexes, test_set_indexes


if __name__ == '__main__':
    # path = '../IMS/2nd_test/2nd_test'
    path = '../IMS/3rd_test/4th_test/txt'
    data_sigma = imsdatasigma(path)
    # datalist = getIMSdatalist(path)
    # trainlist, test_list = ims_dataset_split(datalist)
    # testdataset = [datalist[i] for i in test_list]
    # # datalist[0]
    # # datalist[1]
    # # data = np.loadtxt(datalist[0][0][0])
    # # data1 = data[:, 0]
    # # data2 = data[:, 1]
    # imstrainingdatagenerator = IMSDataGenerator(testdataset)
    # x, y = imstrainingdatagenerator.getitem(3)
    # print(x.shape)
    # imstrainingdatagenerator.getitem(4)
    # xlist, ylist = getdatalist(path, step=3)
    # # print(datalist[5]['value'][0], datalist[5]['age'])
    # data = np.loadtxt(xlist[0][1])
    # y = ylist[0]
    # print(y)
# datadir = getdata(path)





